#!/usr/bin/python3
# -*- coding: utf-8 -*-
import os
import struct
import multiprocessing

video_list_count = 8
base_path = r'/rootfs/media/kasim/DataSet/RecordEx'
out_path = r'/rootfs/media/kasim/DataSet/RecordEx'
video_list_bbox_count_file = r'video_list{}_bbox_count.txt'
video_list_bbox_file = r'video_list{}_bbox.dat'
video_list_size_file = r'video_list{}.txt'

PROCESS_COUNT = video_list_count

min_video_size = 1024*1024*7
max_video_size = 1024*1024*100

low_select_obj_count = 1
high_select_obj_count = 1
min_frame_count = 10

select_class = 'Person'

select_class_label = {
    'Person': 0,
    'Cat': 1,
    'Dog': 2,
    'BabyCar': 3,
    'Face': 4,
}
select_label = select_class_label[select_class]

LOW_THRESHOLDS = [
    0.3,  # Person
    0.3,  # Cat
    0.3,  # Dog
    0.3,  # BabyCar
    0.4,  # Face
]

HIGH_THRESHOLDS = [
    0.6,  # Person
    0.5,  # Cat
    0.5,  # Dog
    0.5,  # BabyCar
    0.6,  # Face
]


def select_proc(file_queue, id):
    try:
        video_size_dict = {}
        with open(os.path.join(base_path, video_list_size_file.format(id)), 'r') as file:
            for video_info in file.readlines():
                video_info = video_info.split(',')
                video_size_dict[video_info[0].strip()] = int(video_info[1].strip())
        print('Proc:{}, Read File Size Finish!'.format(id))

        video_list = []
        with open(os.path.join(base_path, video_list_bbox_count_file.format(id)), 'r') as file:
            for video_info in file.readlines():
                video_info = video_info.split(',')
                # video_list.append([video_info[0].strip(), int(video_info[1].strip())])
                video_name = video_info[0].strip()
                video_size = video_size_dict.get(video_name, 0)
                video_list.append([video_name, int(video_info[1].strip()), video_size])
        print('Proc:{}, Read File BBox Count Finish!'.format(id))

        low_thr = LOW_THRESHOLDS[select_label]
        high_thr = HIGH_THRESHOLDS[select_label]
        file_count = 0
        total_file_count = len(video_list)
        with open(os.path.join(base_path, video_list_bbox_file.format(id)), 'rb') as file:
            for video_name, BBox_count, video_size in video_list:
                if video_size < min_video_size or max_video_size < video_size:
                    bbox_info = file.read(28*BBox_count)
                    if len(bbox_info) < 28*BBox_count:
                        print('Proc:{}, File Count: {}/{}'.format(id, file_count, total_file_count))
                        file_queue.put([None, True, id])
                        return
                    continue
                low_frame_bbox_count = {}
                high_frame_bbox_count = {}
                for j in range(BBox_count):
                    bbox_info = file.read(28)
                    if len(bbox_info) < 28:
                        print('Proc:{}, File Count: {}/{}'.format(id, file_count, total_file_count))
                        file_queue.put([None, True, id])
                        return
                    bbox = struct.unpack('6i1f', bbox_info)
                    label = bbox[1]
                    if select_label != label:
                        continue
                    frame_id = bbox[0]
                    score = bbox[6]
                    if score >= high_thr:
                        if frame_id not in high_frame_bbox_count.keys():
                            high_frame_bbox_count[frame_id] = 1
                        else:
                            high_frame_bbox_count[frame_id] += 1
                    if score >= low_thr:
                        if frame_id not in low_frame_bbox_count.keys():
                            low_frame_bbox_count[frame_id] = 1
                        else:
                            low_frame_bbox_count[frame_id] += 1

                max_low_select_obj_count = 0
                max_low_select_frame_count = 0
                max_low_frame_count = 0
                max_low_select_obj_count_frame_count = 0
                if len(low_frame_bbox_count) > 0:
                    # low_frame_bbox_count记录每一帧中满足低阈值要求的目标数
                    low_frame_bbox_counts = low_frame_bbox_count.values()
                    low_max_bbox_count = max(low_frame_bbox_counts)  # 同一帧里最大目标数
                    if 0 < low_max_bbox_count <= low_select_obj_count:  # 只记录目标数小于等于low_select_obj_count的视频
                        frame_count = 0
                        obj_count_frame_count = 0
                        for bbox_count in low_frame_bbox_counts:
                            if 0 < bbox_count <= low_select_obj_count:  # 记录视频中目标数大于0且小于等于low_select_obj_count的帧数
                                frame_count += 1
                            if bbox_count == low_select_obj_count:  # 记录视频中目标数等于low_select_obj_count的帧数
                                obj_count_frame_count += 1
                            if 0 < bbox_count:
                                max_low_frame_count += 1
                        max_low_select_frame_count = frame_count
                        max_low_select_obj_count_frame_count = obj_count_frame_count
                        max_low_select_obj_count = low_max_bbox_count  # 最大目标数
                max_high_select_obj_count = 0
                max_high_select_frame_count = 0
                max_high_frame_count = 0
                max_high_select_obj_count_frame_count = 0
                if len(high_frame_bbox_count) > 0:
                    # high_frame_bbox_count记录每一帧中满足高阈值要求的目标数
                    high_frame_bbox_counts = high_frame_bbox_count.values()
                    high_max_bbox_count = max(high_frame_bbox_counts)  # 同一帧里最大目标数
                    if high_select_obj_count <= high_max_bbox_count:  # 只记录目标数大于等于high_max_bbox_count的视频
                        frame_count = 0
                        obj_count_frame_count = 0
                        for bbox_count in high_frame_bbox_counts:
                            if high_select_obj_count <= bbox_count:  # 记录视频中目标数大于等于high_select_obj_count的帧数
                                frame_count += 1
                            if bbox_count == high_select_obj_count:  # 记录视频中目标数等于high_select_obj_count的帧数
                                obj_count_frame_count += 1
                            if 0 < bbox_count:
                                max_high_frame_count += 1
                        max_high_select_frame_count = frame_count
                        max_high_select_obj_count_frame_count = obj_count_frame_count
                        max_high_select_obj_count = high_max_bbox_count  # 最大目标数

                if (max_low_select_obj_count > 0) or (max_high_select_obj_count > 0):
                    file_queue.put([(video_name, max_low_select_obj_count, max_high_select_obj_count, max_low_select_frame_count, max_high_select_frame_count, max_low_frame_count, max_high_frame_count, max_low_select_obj_count_frame_count, max_high_select_obj_count_frame_count), False, id])

                file_count += 1
                if file_count % 1000 == 0:
                    print('Proc:{}, File Count: {}/{}'.format(id, file_count, total_file_count))
                    # return
    except Exception as e:
        print('Proc:{}, File Count: {}/{}'.format(id, file_count, total_file_count))
    print('Proc:{}, File Count: {}/{}'.format(id, file_count, total_file_count))
    file_queue.put([None, True, id])


def main():
    if not os.path.exists(out_path):
        os.makedirs(out_path)
    os.system('chmod a+wr {}'.format(out_path))

    file_queue = multiprocessing.Queue()
    select_video_workers = []
    for i in range(PROCESS_COUNT):
        select_video_workers.append(multiprocessing.Process(target=select_proc, args=(file_queue, i,)))

    for i in range(PROCESS_COUNT):
        select_video_workers[i].start()

    low_select_file_list = []
    high_select_file_list = []
    try:
        finish_worker_count = 0
        while True:
            file_info = file_queue.get(block=True)
            if file_info is None:
                break
            video_info, finish, id = file_info
            if finish:
                print('Proc{} finish'.format(id, ))
                finish_worker_count += 1
                if PROCESS_COUNT <= finish_worker_count:
                    break
                continue
            if (video_info is None) or (len(video_info) < 1):
                continue
            video_name, max_low_select_obj_count, max_high_select_obj_count, max_low_select_frame_count, max_high_select_frame_count, max_low_frame_count, max_high_frame_count, max_low_select_obj_count_frame_count, max_high_select_obj_count_frame_count = video_info
            if (max_low_select_obj_count > 0) and (max_low_select_frame_count >= min_frame_count):
                low_select_file_list.append((video_name, max_low_select_obj_count, max_low_select_frame_count))
            if (max_high_select_obj_count > 0) and (max_high_select_frame_count >= min_frame_count):
                high_select_file_list.append((video_name, max_high_select_obj_count, max_high_select_frame_count))
    except Exception as e:
        print(e)

    for i in range(PROCESS_COUNT):
        select_video_workers[i].join()

    # resize_proc(file_queue, 0)

    print('Select Video Workers Finish!')

    print('Read File Queue Finish!')
    file_path = os.path.join(out_path, 'low_select_{}_file_list.txt'.format(select_class))
    with open(file_path, 'w') as file:
        low_select_file_list.sort(key=lambda x: x[0])
        for select_file_info in low_select_file_list:
            select_file_info_str = '{} {} {}\n'.format(select_file_info[0], select_file_info[1], select_file_info[2])
            file.write(select_file_info_str)
    os.system('chmod a+wr {}'.format(file_path))
    print('Write Low Select File {} List {} Finish!'.format(select_class, len(low_select_file_list)))
    file_path = os.path.join(out_path, 'high_select_{}_file_list.txt'.format(select_class))
    with open(file_path, 'w') as file:
        high_select_file_list.sort(key=lambda x: x[0])
        for select_file_info in high_select_file_list:
            select_file_info_str = '{} {} {}\n'.format(select_file_info[0], select_file_info[1], select_file_info[2])
            file.write(select_file_info_str)
    os.system('chmod a+wr {}'.format(file_path))
    print('Write High Select File {} List {} Finish!'.format(select_class, len(high_select_file_list)))
    print('finish!!!!')


if __name__ == '__main__':
    main()
